# -*- coding: utf-8 -*-
# @Function: 统计QA_final目录中Excel文件的问答对数量

import os
import pandas as pd
from loguru import logger

# 设置日志格式
logger.remove()
logger.add(lambda msg: print(msg, flush=True), level="INFO", format="{time:YYYY-MM-DD HH:mm:ss} | {level} | {message}")

# QA_final目录路径
QA_FINAL_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "QA_final")

def count_qa_pairs():
    """统计QA_final目录中各Excel文件的问答对数量"""
    total_count = 0
    file_counts = {}
    
    # 检查目录是否存在
    if not os.path.exists(QA_FINAL_DIR):
        logger.error(f"目录不存在: {QA_FINAL_DIR}")
        return
    
    # 获取目录中的所有Excel文件
    excel_files = [f for f in os.listdir(QA_FINAL_DIR) if f.endswith(".xlsx")]
    
    if not excel_files:
        logger.warning(f"在 {QA_FINAL_DIR} 中未找到Excel文件")
        return
    
    logger.info(f"开始统计 {QA_FINAL_DIR} 中的Excel文件问答对数量...")
    
    # 遍历每个Excel文件并统计行数
    for excel_file in excel_files:
        file_path = os.path.join(QA_FINAL_DIR, excel_file)
        try:
            # 读取Excel文件
            df = pd.read_excel(file_path)
            row_count = len(df)
            
            # 记录文件行数
            file_counts[excel_file] = row_count
            total_count += row_count
            
            logger.info(f"{excel_file}: {row_count} 对问答")
            
        except Exception as e:
            logger.error(f"处理文件 {excel_file} 时出错: {e}")
    
    # 输出总计
    logger.info(f"总计: {total_count} 对问答")
    return file_counts, total_count

if __name__ == "__main__":
    count_qa_pairs()